import numpy as np


def euclidean_distance(x1, x2):
    """
    计算两个向量之间的欧几里得距离
    """
    return np.sqrt(np.sum((x1 - x2) ** 2))


class KMeans:

    def __init__(self, n_clusters=8, max_iter=300, random_state=None, n_init=10, tol=1e-4):
        """
        初始化参数

        参数:
        n_clusters (int): 聚类的簇数，默认值为8
        max_iter (int): 最大迭代次数，防止算法无法收敛一直运行，默认值为300
        random_state (int 或 None): 随机数种子，用于可重复性，若为None则随机生成
        n_init (int): 重复运行K-Means算法的次数，取最优结果，默认值为10
        tol (float): 收敛的容忍度，判断聚类中心是否收敛的阈值，默认值为1e-4
        """
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.random_state = random_state
        self.n_init = n_init
        self.tol = tol
        self.centroids_ = None
        self.labels_ = None
        self.inertia_ = None

    def fit(self, X):
        """
        对输入数据进行聚类

        参数:
        X (numpy.ndarray): 输入的数据，形状为 (n_samples, n_features)，即样本数乘以特征数

        返回:
        self: 返回自身对象，以便可以链式调用（如 kmeans.fit(X).predict(X_new) 等）
        """
        best_inertia = np.inf
        best_centroids = None
        best_labels = None

        for _ in range(self.n_init):
            if self.random_state is not None:
                np.random.seed(self.random_state)
            n_samples, n_features = X.shape
            centroids = X[np.random.choice(n_samples, self.n_clusters, replace=False)]

            for iter_num in range(self.max_iter):
                labels = np.zeros(n_samples)
                for i in range(n_samples):
                    distances = [euclidean_distance(X[i], centroid) for centroid in centroids]
                    labels[i] = np.argmin(distances)

                new_centroids = np.zeros((self.n_clusters, n_features))
                for j in range(self.n_clusters):
                    points_in_cluster = X[labels == j]
                    if len(points_in_cluster) > 0:
                        new_centroids[j] = np.mean(points_in_cluster, axis=0)

                shift = np.linalg.norm(centroids - new_centroids)
                centroids = new_centroids

                if shift < self.tol:
                    break

                inertia = self._calculate_inertia(X, centroids, labels)
                if inertia < best_inertia:
                    best_inertia = inertia
                    best_centroids = centroids
                    best_labels = labels

        self.centroids_ = best_centroids
        self.labels_ = best_labels
        self.inertia_ = best_inertia

        return self

    def predict(self, X):
        """
        预测输入数据所属的簇标签

        参数:
        X (numpy.ndarray): 输入的数据，形状为 (n_samples, n_features)

        返回:
        labels (numpy.ndarray): 每个样本所属的簇标签，形状为 (n_samples,)
        """
        n_samples = X.shape[0]
        labels = np.zeros(n_samples)
        for i in range(n_samples):
            distances = [euclidean_distance(X[i], centroid) for centroid in self.centroids_]
            labels[i] = np.argmin(distances)
        return labels

    def _calculate_inertia(self, X, centroids, labels):
        """
        计算聚类的惯性（所有样本到其所属聚类中心距离的平方和）
        """
        inertia = 0
        for i in range(len(X)):
            inertia += euclidean_distance(X[i], centroids[int(labels[i])]) ** 2
        return inertia